# -*- coding: utf-8 -*-
import json
import re
import datetime
from typing import List, Dict
import requests, os
from parsel import Selector

CACHE_FILE = "asset/calendar_cache.json"

class HolidayCalculator:

    def __init__(self, cache_file: str = CACHE_FILE):
        self.cache_file = cache_file
        self.calendar_cache = self._load_cache()

    def _load_cache(self) -> Dict[str, Dict]:
        """从文件加载缓存数据"""
        if os.path.exists(self.cache_file):
            try:
                with open(self.cache_file, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception as e:
                print(f"⚠️ 缓存文件 {self.cache_file} 加载失败，使用空缓存：{e}")
        return {}

    def _save_cache(self):
        """将缓存数据写入文件（覆盖）"""
        try:
            with open(self.cache_file, "w", encoding="utf-8") as f:
                json.dump(self.calendar_cache, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"❌ 缓存保存失败：{e}")

    def get_calendar(self, year_and_month: str) -> dict:
        """
        获取指定月份的日历数据

        Args:
            year_and_month: 年月字符串，格式 "YYYY-MM"

        Returns:
            日历数据字典
        """
        if year_and_month in self.calendar_cache:
            return self.calendar_cache[year_and_month]

        url = "https://wannianrili.bmcx.com/ajax/"
        params = {"q": year_and_month, "v": "22121303"}
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:64.0) Gecko/20100101 Firefox/64.0"
        }

        try:
            res = requests.get(url, params=params, headers=headers, timeout=10)
            sel = Selector(text=res.text)
            rows = sel.css(".wnrl_riqi")

            data = {}
            for row in rows:
                class_name = row.css("a::attr(class)").extract_first("").strip()
                day = row.css("a::attr(onclick)").extract_first("").strip()
                comment = row.css(".wnrl_td_bzl::text").extract_first("").strip()

                ret = re.search(r"\d{4}-\d{2}-\d{2}", day)
                if ret:
                    day = ret.group(0)
                    data[day] = {"class_name": class_name, "comment": comment}

            # 更新缓存并持久化
            self.calendar_cache[year_and_month] = data
            self._save_cache()
            return data

        except Exception as e:
            print(f"获取日历数据失败: {e}")
            return {}

    def get_day_item(self, day: str) -> dict:
        """
        获取指定日期的日历项

        Args:
            day: 日期字符串，格式 "YYYY-MM-DD"

        Returns:
            日历项字典
        """
        year_month = "-".join(day.split("-")[:2])
        calendar = self.get_calendar(year_month)
        return calendar.get(day, {})

    def is_workday(self, day: str) -> bool:
        """
        判断是否为工作日

        Args:
            day: 日期字符串，格式 "YYYY-MM-DD"

        Returns:
            True如果是工作日，否则False
        """
        # 工作日class列表
        workday_class_list = ["", "wnrl_riqi_ban"]

        day_item = self.get_day_item(day)
        if day_item:
            return day_item.get("class_name") in workday_class_list
        return False

    def is_holiday(self, day: str) -> bool:
        """
        判断是否为节假日

        Args:
            day: 日期字符串，格式 "YYYY-MM-DD"

        Returns:
            True如果是节假日，否则False
        """
        # 节假日class列表
        holiday_class_list = ["wnrl_riqi_xiu", "wnrl_riqi_mo"]

        day_item = self.get_day_item(day)
        if day_item:
            return day_item.get("class_name") in holiday_class_list
        return False

    def get_holidays(self, start_time: str, end_time: str) -> List[str]:
        """
        获取指定时间段内的节假日日期

        Args:
            start_time: 开始时间，格式 "YYYY-MM-DD"
            end_time: 结束时间，格式 "YYYY-MM-DD"

        Returns:
            非节假日日期列表，格式 "YYYY-MM-DD"
        """
        # 验证日期格式
        try:
            start_date = datetime.datetime.strptime(start_time, "%Y-%m-%d")
            end_date = datetime.datetime.strptime(end_time, "%Y-%m-%d")
        except ValueError:
            raise ValueError("日期格式错误，请使用YYYY-MM-DD格式")

        if start_date > end_date:
            raise ValueError("开始时间不能晚于结束时间")

        holidays = []
        current_date = start_date

        # 遍历日期范围
        while current_date <= end_date:
            current_day_str = current_date.strftime("%Y-%m-%d")

            # 如果是非节假日，添加到结果列表
            if self.is_holiday(current_day_str):
                holidays.append(current_day_str)

            current_date += datetime.timedelta(days=1)

        return holidays

    def get_workdays(self, start_time: str, end_time: str) -> List[str]:
        """
        获取指定时间段内的工作日（非节假日且非周末）

        Args:
            start_time: 开始时间，格式 "YYYY-MM-DD"
            end_time: 结束时间，格式 "YYYY-MM-DD"

        Returns:
            工作日日期列表，格式 "YYYY-MM-DD"
        """
        try:
            start_date = datetime.datetime.strptime(start_time, "%Y-%m-%d")
            end_date = datetime.datetime.strptime(end_time, "%Y-%m-%d")
        except ValueError:
            raise ValueError("日期格式错误，请使用YYYY-MM-DD格式")

        if start_date > end_date:
            raise ValueError("开始时间不能晚于结束时间")

        workdays = []
        current_date = start_date

        while current_date <= end_date:
            current_day_str = current_date.strftime("%Y-%m-%d")

            # 如果是工作日（非节假日且非周末）
            if self.is_workday(current_day_str):
                workdays.append(current_day_str)

            current_date += datetime.timedelta(days=1)

        return workdays

    def get_date_range_months(self, start_time: str, end_time: str) -> List[str]:
        """
        获取日期范围内涉及的所有月份

        Args:
            start_time: 开始时间
            end_time: 结束时间

        Returns:
            月份列表，格式 ["YYYY-MM", ...]
        """
        start_date = datetime.datetime.strptime(start_time, "%Y-%m-%d")
        end_date = datetime.datetime.strptime(end_time, "%Y-%m-%d")

        months = []
        current_date = start_date

        while current_date <= end_date:
            month_str = current_date.strftime("%Y-%m")
            if month_str not in months:
                months.append(month_str)

            # 跳到下个月的第一天
            if current_date.month == 12:
                current_date = datetime.datetime(current_date.year + 1, 1, 1)
            else:
                current_date = datetime.datetime(
                    current_date.year, current_date.month + 1, 1
                )

        return months

    def preload_calendar_data(self, start_time: str, end_time: str):
        """
        预加载日期范围内的日历数据

        Args:
            start_time: 开始时间
            end_time: 结束时间
        """
        months = self.get_date_range_months(start_time, end_time)
        for month in months:
            self.get_calendar(month)  # 这会自动缓存数据


# 使用示例
if __name__ == "__main__":
    calculator = HolidayCalculator()

    # 示例：计算2024年5月的节假日/工作日
    start_time = "2024-05-01"
    end_time = "2024-05-31"

    try:
        # 预加载数据（可选，可以提高性能）
        calculator.preload_calendar_data(start_time, end_time)

        # 获取节假日
        holidays = calculator.get_holidays(start_time, end_time)

        print(f"{start_time} 到 {end_time} 的节假日共 {len(holidays)} 天:")
        for i, day in enumerate(holidays, 1):
            day_item = calculator.get_day_item(day)
            comment = day_item.get("comment", "")
            print(f"{i}. {day} {comment}")

        print("\n" + "=" * 50 + "\n")

        # 获取工作日
        workdays = calculator.get_workdays(start_time, end_time)
        print(f"{start_time} 到 {end_time} 的工作日共 {len(workdays)} 天:")
        for i, day in enumerate(workdays, 1):
            day_item = calculator.get_day_item(day)
            comment = day_item.get("comment", "")
            print(f"{i}. {day} {comment}")

    except Exception as e:
        print(f"计算错误: {e}")
